#include <msg.h>
#include <utils.h>
#include <main.h>

#include <game.h>

typedef void (*fill_fn)(struct game_data *buf);

struct name_fn_pair {
    const wchar_t *name;
    fill_fn fill;
};

const struct name_fn_pair GAMES[] = {
    { L"BH3",                   &hi3_fill_data  },
    { L"StarRail",              &hsr_fill_data  },
    { L"Client-Win64-Shipping", &wuwa_fill_data }
};

void game_detect(struct game_data *buf) {
    wchar_t exe_path[MAX_PATH];
    GetModuleFileNameW(NULL, exe_path, MAX_PATH);

    // Leave only the basename
    wchar_t *exe_name = wcsrchr(exe_path, L'\\') + 1;

    // Cut off extension (.exe)
    wchar_t *extension_dot = wcsrchr(exe_name, L'.');
    if (extension_dot != NULL) {
        *extension_dot = L'\0';
    }

    for (size_t i = 0; i < UTILS_COUNT(GAMES); i++) {
        if (wcsicmp(exe_name, GAMES[i].name) == 0) {
            GAMES[i].fill(buf);
            return;
        }
    }

    if (!utils_env_enabled("JADEITE_ALLOW_UNKNOWN")) {
        msg_err_w(L"Unknown game: \"%ls\"\n\nSet JADEITE_ALLOW_UNKNOWN=1 to try continuing anyway in skip mode (only Unity games are supported)", exe_name);
    }

    buf->shell_sub = SHELL_SKIP;
    buf->base_sub = BASE_NONE;
}

void *game_generate_unity_loader() {
    char template[] = {
        0x48, 0xb9, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  // mov rcx, instance   ; instance      +0x2
        0x48, 0xc7, 0xc2, 0x00, 0x00, 0x00, 0x00,                    // mov rdx, NULL       ; prevInstance
        0x49, 0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  // mov r8,  cmdline    ; cmdline       +0x13
        0x49, 0xc7, 0xc1, 0x01, 0x00, 0x00, 0x00,                    // mov r9,  SW_NORMAL  ; show

        0xFF, 0x25, 0x00, 0x00, 0x00, 0x00,                          // jmp [$+6]           ; UnityMain()
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00               // dq UnityMain        ; fn addr       +0x28
    };

    HMODULE unityplayer = LoadLibraryA("UnityPlayer.dll");
    if (!unityplayer) {
        msg_err_a("Failed to load UnityPlayer.dll");
    }

    void *unitymain = GetProcAddress(unityplayer, "UnityMain");
    if (!unitymain) {
        msg_err_a("Failed to find UnityMain, UnityPlayer.dll is corrupt or obfuscated");
    }

    *((HMODULE*)(template + 0x2)) = GetModuleHandleA(NULL);
    *((void**)(template + 0x13)) = GetCommandLineW();
    *((void**)(template + 0x28)) = unitymain;

    void *page = VirtualAlloc(NULL, sizeof(template), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    memcpy(page, template, sizeof(template));

    return page;
}

void *game_find_msvc_entry_point() {
    // calls that appear in canary init function, in order
    char *ENTRY_FUNCTION_NAMES[4] = {
        "GetSystemTimeAsFileTime",
        "GetCurrentThreadId",
        "GetCurrentProcessId",
        "QueryPerformanceCounter"
    };

    void *entry_function_ptrs[4];
    HMODULE kernel32 = GetModuleHandleA("kernel32.dll");
    for (size_t i = 0; i < 4; i++) {
        entry_function_ptrs[i] = GetProcAddress(kernel32, ENTRY_FUNCTION_NAMES[i]);
    }

    // get exe headers
    char *exe = (char*)GetModuleHandleA(NULL);
    IMAGE_DOS_HEADER *dos_header = (IMAGE_DOS_HEADER*)exe;
    IMAGE_NT_HEADERS64 *nt_headers = (IMAGE_NT_HEADERS64*)(exe + dos_header->e_lfanew);

    // get runtime function array
    IMAGE_DATA_DIRECTORY rfa_dir = nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
    size_t count = rfa_dir.Size / sizeof(_IMAGE_RUNTIME_FUNCTION_ENTRY);
    _IMAGE_RUNTIME_FUNCTION_ENTRY *rfa = (_IMAGE_RUNTIME_FUNCTION_ENTRY*)(exe + rfa_dir.VirtualAddress);

    // get rdata section
    IMAGE_SECTION_HEADER *rdata = ((IMAGE_SECTION_HEADER*)(nt_headers + 1)) + 1;
    char *rdata_ptr = exe + rdata->VirtualAddress;

    // try to find the canary init function
    void *canary_init = NULL;
    for (size_t i = 0; i < count; i++) {
        size_t seq = 0;
        for (char *ptr = exe + rfa[i].BeginAddress; ptr < exe + rfa[i].EndAddress; ptr++) {
            // call
            if (*((uint16_t*)ptr) != 0x15FF) {
                continue;
            }

            void **import_addr = (void**)(ptr + *((int32_t*)(ptr + 2)) + 6);
            if ((char*)import_addr < rdata_ptr || (char*)import_addr > rdata_ptr + rdata->SizeOfRawData) {
                continue;
            }

            if (*import_addr != entry_function_ptrs[seq]) {
                continue;
            }

            seq++;
            if (seq == 4) {
                break;
            }

            ptr += 5; // skip call, +1 from for
        }
        if (seq == 4) {
            canary_init = exe + rfa[i].BeginAddress;
            break;
        }
    }
    if (!canary_init) {
        msg_err_a("Failed to find canary init function. " ISSUE_SUFFIX);
    }

    // use it to find the original entry point
    char template[] = {
        0x48, 0x83, 0xEC, 0x28,       // sub rsp, 28
        0xE8, 0x00, 0x00, 0x00, 0x00, // call canary init (offset +5, rip +9)
        0x48, 0x83, 0xC4, 0x28,       // add rsp, 28
        0xE9                          // jmp
    };

    void *entry_point = NULL;
    for (size_t i = 0; i < count; i++) {
        char *ptr = exe + rfa[i].BeginAddress;

        *((int32_t*)(template + 5)) = (int32_t)((ptrdiff_t)canary_init - (ptrdiff_t)(ptr + 9));
        if (memcmp(ptr, template, sizeof(template)) == 0) {
            entry_point = ptr;
            break;
        }
    }
    if (!entry_point) {
        msg_err_a("Failed to find entry point. " ISSUE_SUFFIX);
    }

    return entry_point;
}
